# 整合代码：招聘一点通

import streamlit as st
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import json
from streamlit_echarts import st_echarts
from streamlit_echarts import JsCode
from getdata import get_count, load_data
import matplotlib
import pickle
from ast import literal_eval
from zhipuai import ZhipuAI
import plotly.graph_objects as go
import plotly.express as px
from wordcloud import WordCloud
import os

# 初始化登录状态
if "is_logged_in" not in st.session_state:
    st.session_state.is_logged_in = False

# 初始化 st.session_state.api_key_input
if "api_key_input" not in st.session_state:
    st.session_state.api_key_input = ""

# 设置Matplotlib的全局字体为Microsoft YaHei
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei']
matplotlib.rcParams['axes.unicode_minus'] = False

# 登录函数
def login():
    st.title("登录页面")
    username = st.text_input("用户名")
    password = st.text_input("密码", type="password")
    if st.button("登录"):
        if username == "users" and password == "123":
            st.session_state.is_logged_in = True
            st.success("登录成功！")
            # 修改为 st.rerun()
            st.rerun()
        else:
            st.error("用户名或密码错误，请重试。")

# 如果未登录，显示登录页面
if not st.session_state.is_logged_in:
    login()
else:
    # -------------------------- AI机器人对话功能 --------------------------

    def load_data():
        df = pd.read_csv("拉勾网2023招聘数据.csv")

        def parse_salary(salary_str):
            if pd.isna(salary_str) or '面议' in salary_str:
                return (0, 0)
            try:
                low, high = salary_str.replace('k', '000').split('-')
                return (int(low), int(high))
            except:
                return (0, 0)

        df['salary_range'] = df['salary'].apply(parse_salary)

        def parse_skills(skills_str):
            if pd.isna(skills_str):
                return []
            try:
                return literal_eval(skills_str)
            except:
                return []

        df['skillLables'] = df['skillLables'].apply(parse_skills)
        return df


    df = load_data()


    def count_positions(city, position_name):
        filtered = df[(df['city'] == city) & (df['positionName'].str.contains(position_name))]
        position_counts = filtered['positionName'].value_counts().reset_index()

        city_data = df[df['city'] == city]
        position_counts_city = city_data['positionName'].value_counts().head(5)
        fig_bar = px.bar(position_counts_city, x=position_counts_city.index, y=position_counts_city.values,
                         title=f"{city}职位分布")
        fig_pie = px.pie(position_counts_city, names=position_counts_city.index, values=position_counts_city.values,
                         title=f"{city}职位占比")
        fig_radar = go.Figure(data=go.Scatterpolar(r=position_counts_city.values, theta=position_counts_city.index,
                                                   fill='toself'))
        fig_radar.update_layout(title=f"{city}职位分布雷达图")

        current_charts = [{"type": "bar", "fig": fig_bar}, {"type": "pie", "fig": fig_pie},
                          {"type": "radar", "fig": fig_radar}]
        return len(filtered), fig_bar


    def average_salary(position_name):
        filtered = df[df['positionName'].str.contains(position_name)]
        salaries = [(low + high) // 2 for low, high in filtered['salary_range'] if (low + high) > 0]
        return int(sum(salaries) / len(salaries)) if salaries else 0


    def max_salary(city, position_name):
        filtered = df[(df['city'] == city) & (df['positionName'].str.contains(position_name))]
        return max([high for low, high in filtered['salary_range']], default=0)


    def load_apriori(filepath):
        with open(filepath, "rb") as f:
            return pickle.load(f)


    def load_skill_pool(skill_csv_path):
        df_skill = pd.read_csv(skill_csv_path)
        all_skills = []
        for skills_str in df_skill['skill_list']:
            skills = literal_eval(skills_str) if isinstance(skills_str, str) else []
            all_skills.extend(skills)
        return list(set(all_skills))


    apriori_data = load_apriori("apriori.bin")
    SKILL_POOL = load_skill_pool("skill_list.csv")


    def get_associated_skills(target_skill, min_support=0.01):
        target_skill = target_skill.lower()
        filtered_rules = apriori_data["rules"][apriori_data["rules"]["antecedents"].apply(lambda x: target_skill in x)]
        raw_skills = list(
            filtered_rules[filtered_rules["support"] >= min_support]["consequents"].explode().unique())
        return [s for s in raw_skills if s in SKILL_POOL]


    def get_position_skills(position_name):
        position_skills = df[df["positionName"].str.contains(position_name)]["skillLables"]
        if len(position_skills) == 0:
            return ["暂无该岗位的技能数据"]
        freq = pd.Series(
            [skill for sublist in position_skills.dropna() for skill in sublist]).value_counts().head(5)
        return freq.index.tolist() if not freq.empty else ["暂无高频技能数据"]


    def get_suitable_positions(skill):
        skill = skill.lower()
        suitable_positions = df[df["skillLables"].apply(lambda x: skill in x)]["positionName"]
        return suitable_positions.value_counts().head(5).index.tolist()


    def get_city_position_distribution(city):
        city_data = df[df['city'] == city]
        position_counts = city_data['positionName'].value_counts().head(5)
        return position_counts


    def get_position_avg_salary_by_city(position_name, top_n=5):
        city_salary_avg = {}
        for city in df['city'].unique():
            city_data = df[(df['city'] == city) & (df['positionName'].str.contains(position_name))]
            if not city_data.empty:
                salaries = [(low + high) // 2 for low, high in city_data['salary_range'] if (low + high) > 0]
                city_salary_avg[city] = int(sum(salaries) / len(salaries)) if salaries else 0
        sorted_cities = sorted(city_salary_avg.items(), key=lambda x: x[1], reverse=True)[:top_n]
        return pd.DataFrame(sorted_cities, columns=['城市', '平均薪资'])


    def get_position_max_salary_by_city(position_name, top_n=5):
        city_salary_max = {}
        for city in df['city'].unique():
            city_data = df[(df['city'] == city) & (df['positionName'].str.contains(position_name))]
            if not city_data.empty:
                max_salaries = [high for low, high in city_data['salary_range']]
                city_salary_max[city] = max(max_salaries) if max_salaries else 0
        sorted_cities = sorted(city_salary_max.items(), key=lambda x: x[1], reverse=True)[:top_n]
        return pd.DataFrame(sorted_cities, columns=['城市', '最高薪资'])


    def get_city_position_count(city, position):
        count = len(df[(df['city'] == city) & (df['positionName'] == position)])
        return count


    # 新增：获取指定城市、学历下top8热门岗位的平均薪资
    def get_edu_salary_top8(city, education):
        """获取指定城市、学历下top8热门岗位的平均薪资"""
        # 筛选数据：城市+学历
        filtered = df[(df['city'] == city) & (df['education'] == education)]

        # 计算热门岗位（按岗位出现次数排序取前8）
        top_positions = filtered['positionName'].value_counts().head(8).index.tolist()
        if not top_positions:
            return pd.DataFrame(columns=['岗位名称', '平均薪资'])

        # 计算每个岗位的平均薪资（薪资处理逻辑同原代码）
        salary_data = []
        for pos in top_positions:
            pos_data = filtered[filtered['positionName'] == pos]
            salaries = [(int(low) * 1000 + int(high) * 1000) // 2  # 将"15k-30k"转换为15000-30000，取中间值
                        for low, high in pos_data['salary'].str.extract(r'(\d+)k-(\d+)k').dropna().values]
            avg_salary = int(sum(salaries) / len(salaries)) if salaries else 0
            salary_data.append({'岗位名称': pos, '平均薪资': avg_salary})

        return pd.DataFrame(salary_data)


    # -------------------------- Streamlit界面配置 --------------------------

    st.set_page_config(page_title="招聘一点通", page_icon="🔍")

    # 会话历史初始化
    if "messages" not in st.session_state:
        st.session_state.messages = [{
            "role": "assistant",
            "content": "你好！我是招聘助手，可查询岗位<=>技术相关信息;支持的提问类型：           1.擅长的技术是xxx，推荐关联技术    2.掌握技术为xxx适合的岗位有   3.从事xxx岗位需要的技术",
            "charts": [],
            "used_function": "",
            "function_result": "",
            "llm_answer": ""
        }]

    # 确保 api_configured 被初始化
    if "api_configured" not in st.session_state:
        st.session_state.api_configured = False

    # 左侧功能选择栏
    with st.sidebar:
        st.title("API 调用")
        def update_api_status():
            st.session_state.api_configured = bool(st.session_state.api_key_input)

        api_key = st.text_input("输入Token", type="password", key="api_key_input",
                                on_change=update_api_status)
        if st.session_state.api_configured:
            st.success("API Token已经配置", icon='✅')
        model = st.selectbox("选择模型", ["glm-4", "glm-3-turbo"])
        temperature = st.slider("Temperature", 0.0, 2.0, 0.8)

        st.title("功能选择")
        selected_mode = st.selectbox("选择功能", [
            "不同城市->学历->薪资功能",
            "岗位查询功能",
            "AI机器人对话"
        ])

        if selected_mode == "岗位查询功能":
            city = st.selectbox("选择城市", df['city'].unique())
            position = st.selectbox("选择岗位", df['positionName'].unique())
            if st.button("查询"):
                func_name = "get_city_position_count"
                func_result = get_city_position_count(city, position)
                if st.session_state.api_configured:
                    client = ZhipuAI(api_key=st.session_state.api_key_input)
                    llm_prompt = f"在{city}，{position}岗位的数量是{func_result}个，请用自然语言描述这个结果。"
                    final_response = client.chat.completions.create(
                        model=model,
                        messages=[{"role": "user", "content": llm_prompt}]
                    )
                    llm_answer = final_response.choices[0].message.content
                else:
                    llm_answer = f"在{city}，{position}岗位的数量是{func_result}个。"
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": llm_answer,
                    "charts": [],
                    "used_function": func_name,
                    "function_result": str(func_result),
                    "llm_answer": llm_answer
                })
        elif selected_mode == "AI机器人对话":
            st.title("招聘智能聊天机器人")

            def clear_history():
                st.session_state.messages = [{
                    "role": "assistant",
                    "content": "你好！我是招聘助手，可查询岗位<=>技术相关信息;支持的提问类型：           1.擅长的技术是xxx，推荐关联技术    2.掌握技术为xxx适合的岗位有   3.从事xxx岗位需要的技术",
                    "charts": [],
                    "used_function": "",
                    "function_result": "",
                    "llm_answer": ""
                }]

            st.button("清空聊天记录", on_click=clear_history)

    # 根据选择的功能加载不同的内容
    if selected_mode == "岗位查询功能":
        st.title("岗位查询结果")
        for msg in st.session_state.messages[-1:]:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])
                if msg["used_function"]:
                    st.markdown(f"**使用的函数名：{msg['used_function']}**")
                if msg["function_result"]:
                    st.markdown(f"**函数原始返回值：{msg['function_result']}**")
                if msg["llm_answer"]:
                    st.markdown(f"**LLM自然语言回答：{msg['llm_answer']}**")
                # 可视化展示
                position_counts = df[(df['city'] == city) & (df['positionName'] == position)].groupby(
                    ['city', 'positionName']).size().reset_index(name='count')
                fig = px.bar(position_counts, x='positionName', y='count',
                             title=f'{city}的{position}岗位数量')
                st.plotly_chart(fig)
    elif selected_mode == "不同城市->学历->薪资功能":
        st.title("不同城市-学历-热门岗位薪资分析")
        col1, col2 = st.columns([1, 3])  # 创建两列布局
        with col1:
            # 保留左侧的选择框
            city = st.selectbox("选择城市", df['city'].unique())
        with col2:
            # 将学历选项放在右侧最上面
            education = st.selectbox("选择学历", df['education'].unique())
            if st.button("查询学历-薪资数据", key="edu_salary_btn"):
                # 1. 调用函数获取top8岗位薪资数据
                func_name = "get_edu_salary_top8"
                func_result = get_edu_salary_top8(city, education)

                if st.session_state.api_configured:
                    client = ZhipuAI(api_key=st.session_state.api_key_input)
                    if func_result.empty:
                        llm_prompt = f"在{city}市，未找到{education}学历的热门岗位薪资数据，请用自然语言描述这个结果。"
                    else:
                        data_str = ""
                        for _, row in func_result.iterrows():
                            data_str += f"{row['岗位名称']}岗位的平均薪资为{row['平均薪资']}元/月。\n"
                        llm_prompt = f"以下是{city}市{education}学历的前8热门岗位平均薪资情况：\n{data_str}请用自然语言描述这个结果。"
                    final_response = client.chat.completions.create(
                        model=model,
                        messages=[{"role": "user", "content": llm_prompt}]
                    )
                    llm_answer = final_response.choices[0].message.content
                else:
                    if func_result.empty:
                        llm_answer = f"在{city}市，未找到{education}学历的热门岗位薪资数据。"
                    else:
                        llm_answer = f"以下是{city}市{education}学历的前8热门岗位平均薪资情况：\n"
                        for _, row in func_result.iterrows():
                            llm_answer += f"{row['岗位名称']}岗位的平均薪资为{row['平均薪资']}元/月。\n"

                # 3. 可视化（条形图）
                fig = px.bar(
                    func_result,
                    x='岗位名称',
                    y='平均薪资',
                    title=f"{city}市{education}学历前8热门岗位平均薪资",
                    text='平均薪资',
                    color='岗位名称'
                )
                fig.update_layout(yaxis_title="平均薪资（元/月）")

                # 4. 记录会话历史并显示结果
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": llm_answer,
                    "charts": [{"type": "bar", "fig": fig}],
                    "used_function": func_name,
                    "function_result": str(func_result),
                    "llm_answer": llm_answer
                })

                # 输出结果
                with st.chat_message("assistant"):
                    st.markdown(f"**使用的函数名：{func_name}**")
                    st.markdown(f"**函数原始返回值：**\n{func_result}")
                    st.markdown(f"**自然语言回答：**\n{llm_answer}")
                    st.plotly_chart(fig, use_container_width=True)
    elif selected_mode == "AI机器人对话":
        # 显示所有历史对话
        for msg in st.session_state.messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])
                if msg["used_function"]:
                    st.markdown(f"**使用的函数：{msg['used_function']}**")
                if msg["function_result"]:
                    st.markdown(f"**函数返回值：{msg['function_result']}**")
                if msg["llm_answer"]:
                    st.markdown(f"**LLM自然语言回答：{msg['llm_answer']}**")
                for chart in msg.get("charts", []):
                    if isinstance(chart, dict) and "type" in chart:
                        if chart["type"] == "wordcloud":
                            wc = WordCloud(width=800, height=400, background_color='white').generate_from_frequencies(
                                chart["data"])
                            fig, ax = plt.subplots()
                            ax.imshow(wc, interpolation='bilinear')
                            ax.axis("off")
                            st.pyplot(fig)
                        elif chart["type"] == "bar":
                            st.plotly_chart(chart["fig"], use_container_width=True)
                        elif chart["type"] == "pie":
                            st.plotly_chart(chart["fig"], use_container_width=True)
                        elif chart["type"] == "radar":
                            st.plotly_chart(chart["fig"], use_container_width=True)
                        elif chart["type"] == "table":
                            st.dataframe(chart["data"])

        # 用户输入和处理
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "recommend_associated_skills",
                    "description": "根据用户擅长的技术，推荐关联的其他技术，用1句简短的自然语言回答",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "skill": {
                                "type": "string",
                                "description": "用户擅长的技术名称（如：java）"
                            }
                        },
                        "required": ["skill"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "get_position_required_skills",
                    "description": "查询某岗位需要掌握的技术，用1句简短的自然语言回答",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "position": {
                                "type": "string",
                                "description": "目标岗位名称（如：算法工程师）"
                            }
                        },
                        "required": ["position"]
                    }
                }
            },
            {
                "type": "function",
                "function": {
                    "name": "find_suitable_positions",
                    "description": "根据用户掌握的技术，推荐适合的岗位，用1句简短的自然语言回答",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "skill": {
                                "type": "string",
                                "description": "用户掌握的技术名称（如：python）"
                            }
                        },
                        "required": ["skill"]
                    }
                }
            }
        ]
        if prompt := st.chat_input("输入你的问题（如：从事java工程师需要哪些技术？）"):
            st.session_state.messages.append(
                {"role": "user", "content": prompt, "charts": [], "used_function": "",
                 "function_result": "", "llm_answer": ""})
            with st.chat_message("user"):
                st.markdown(prompt)

            with st.chat_message("assistant"):
                with st.spinner("思考中..."):
                    if not st.session_state.api_configured:
                        error_msg = "请先配置API Key"
                        st.error(error_msg, icon='⚠️')
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": error_msg,
                            "charts": [],
                            "used_function": "",
                            "function_result": "",
                            "llm_answer": ""
                        })
                    else:
                        client = ZhipuAI(api_key=st.session_state.api_key_input)
                        response = client.chat.completions.create(
                            model=model,
                            messages=[{"role": "user", "content": prompt}],
                            tools=tools,
                            tool_choice="auto"
                        )
                        message = response.choices[0].message
                        tool_calls = message.tool_calls if hasattr(message, "tool_calls") else []

                        output = ""
                        current_charts = []

                        if tool_calls:
                            tool_call = tool_calls[0]
                            function_name = tool_call.function.name
                            params = json.loads(tool_call.function.arguments)

                            # 调用相应的工具函数
                            if function_name == "recommend_associated_skills":
                                func_result = get_associated_skills(params["skill"])
                                wordcloud_data = pd.Series(func_result).value_counts().to_dict()
                                current_charts = [{"type": "wordcloud", "data": wordcloud_data}]
                                llm_prompt = f"用户擅长的技术是{params['skill']}，推荐的关联技术是{func_result}"
                            elif function_name == "find_suitable_positions":
                                func_result = get_suitable_positions(params["skill"])
                                suitable_positions_df = pd.DataFrame({"适合的岗位": func_result})
                                current_charts = [{"type": "table", "data": suitable_positions_df}]
                                llm_prompt = f"掌握技术为{params['skill']}适合的岗位有{func_result}"
                            elif function_name == "get_position_required_skills":
                                func_result = get_position_skills(params["position"])
                                skill_counts = pd.Series(func_result).value_counts().reset_index(
                                    name='count')
                                current_charts = [
                                    {"type": "table",
                                     "data": pd.DataFrame({"核心技能": skill_counts['index'],
                                                           "出现频次": skill_counts['count']})}]
                                llm_prompt = f"从事{params['position']}岗位需要的技术是{func_result}"

                            # 调用LLM生成自然语言回答
                            final_response = client.chat.completions.create(
                                model=model,
                                messages=[{"role": "user", "content": llm_prompt}]
                            )
                            llm_answer = final_response.choices[0].message.content

                            # 记录对话历史
                            current_message = {
                                "role": "assistant",
                                "content": llm_answer,
                                "charts": current_charts,
                                "used_function": function_name,
                                "function_result": str(func_result),
                                "llm_answer": llm_answer
                            }
                            st.session_state.messages.append(current_message)

                            # 输出调试信息和图表
                            st.markdown(f"**使用的函数：{function_name}**", unsafe_allow_html=True)
                            st.markdown(f"**函数返回值：{func_result}**", unsafe_allow_html=True)
                            st.markdown(f"**LLM自然语言回答：{llm_answer}**", unsafe_allow_html=True)

                            # 渲染图表
                            for chart in current_charts:
                                if chart["type"] == "bar":
                                    st.plotly_chart(chart["fig"], use_container_width=True)
                                elif chart["type"] == "table":
                                    st.dataframe(chart["data"])
                                elif chart["type"] == "radar":
                                    st.plotly_chart(chart["fig"], use_container_width=True)
                                elif chart["type"] == "wordcloud":
                                    wc = WordCloud(width=800, height=400,
                                                   background_color='white').generate_from_frequencies(
                                        chart["data"])
                                    fig, ax = plt.subplots()
                                    ax.imshow(wc, interpolation='bilinear')
                                    ax.axis("off")
                                    st.pyplot(fig)
                                elif chart["type"] == "pie":
                                    st.plotly_chart(chart["fig"], use_container_width=True)
